import unittest
import references_classifier

references = ["那有一只猪", "哪有一只猪", "他说你是猪", "我说你才是猪", "我说我才是猪"]

class TestRefClassifier(unittest.TestCase):
    def setUp(self):
        self._classifier =  references_classifier.RefClassifier(0.5)

    def test_ref_classify(self):
        self._classifier.ref_classify("谁还不是只猪呢")

    def test_refs_classify(self):
        self._classifier.refs_classify(references)

    def test_export(self):
        self._classifier.refs_classify(references)
        self._classifier.export("ref_count.xlsx")

    def test_compare(self):
        print(self._classifier._compare(references[0], references[1]))
        
if __name__ == "__main__":
    unittest.main()